#include "llvm/Module.h"
#include "llvm/LLVMContext.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Instructions.h"
#include "llvm/Constants.h"
#include "llvm/BasicBlock.h"
#include "llvm/Pass.h"
#include "llvm/Support/CFG.h"
#include "llvm/Transforms/Utils/BasicInliner.h"
#include "llvm/Target/TargetData.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
#include "llvm/Analysis/Dominators.h"
#include "llvm/Analysis/DebugInfo.h"
#include "llvm/Transforms/Utils/SSAUpdater.h"
#include "llvm/Analysis/PostDominators.h"

#include <set>
#include <map>
#include <fstream>
#include <ext/hash_map>

#include "../../../../include/util.h"
#include "../../../../idm/id.h"
#include "omit-branch.h"
#include "../sync-func/sync-func.h"

using namespace llvm;
using namespace std;
using namespace __gnu_cxx;

namespace defens {

	void OmitBranch::getAnalysisUsage(AnalysisUsage &AU) const {
		AU.setPreservesAll();
		AU.addRequired<SyncFunc>();
		AU.addRequired<PostDominatorTree>();
		ModulePass::getAnalysisUsage(AU);
	}

	bool OmitBranch::runOnModule(Module &M) {
		SF = &getAnalysis<SyncFunc>();
		brInstOmitted.clear();
		simple_test(M);
		return false;
	}

	/* Test whether the first branch of main function can be omitted. */
	void OmitBranch::simple_test(Module &M) {
		BranchInst *ins = NULL;
		forallfunc(M, fi) {
			//if (fi->getName() != "main")
			//	continue;
			for (Function::iterator bi = fi->begin(); bi != fi->end(); ++bi) {
				for (BasicBlock::iterator ii = bi->begin(); ii != bi->end(); ++ii) {
					if (isa<BranchInst>(ii)) {
						ins = dyn_cast<BranchInst>(ii);
						/*if (ins->getParent()->getParent()->getNameStr() == "open_entity") {
							fprintf(stderr, "BrInst %p: %d\n", (void *)ins, omit(ins));
							errs() << "BrInst " << *ins << ", at BB " << ins->getParent()->getNameStr() 
								<< ", at Func " << ins->getParent()->getParent()->getNameStr() << ".\n";
							fprintf(stderr, "-----------------------------------------------\n\n\n");
						}*/
						if (omit(ins)) {
							brInstOmitted[(long)ins] = 1;
						} else {
							brInstOmitted[(long)ins] = 0;
						}
						//break;
					}
				}
				//if (ins)
				//	break;
			}
			//if (ins)
			//	break;
		}
		if (!ins) {
			fprintf(stderr, "No branch instructions in this module\n");
			return;
		}

		// Heming Test.
		/*Function *atomF = M.getFunction("apr_atomic_inc");
		assert(atomF);
		if (SF->contains_sync_operation(atomF)) {
			fprintf(stderr, "apr_atomic_inc contains sync operations\n");
			exit(0);
		} else {
			exit(1);
		}*/
		//fprintf(stderr, "%d\n", omit(ins));
	}

	void OmitBranch::print(raw_ostream &O, const Module *M) const {
	}

	void OmitBranch::dfs(BasicBlock *x, BasicBlock *sink) {
		// cerr << "DFS: " << x->getNameStr() << endl;
		if (visited[x])
			return;
		// Stop at the sink -- the post dominator of the branch
		if (x == sink)
			return;
		visited[x] = true;
		for (succ_iterator it = succ_begin(x); it != succ_end(x); ++it) {
			BasicBlock *y = *it;
			dfs(y, sink);
		}
	}

	/* Copied from LLVM SVN */
	BasicBlock *find_nearest_common_post_dominator(PostDominatorTree &PDT, BasicBlock *A, BasicBlock *B) {
		// If B dominates A then B is nearest common dominator.
		if (PDT.dominates(B, A))
			return B;

		// If A dominates B then A is nearest common dominator.
		if (PDT.dominates(A, B))
			return A;

		DomTreeNodeBase<BasicBlock> *NodeA = PDT.getNode(A);
		DomTreeNodeBase<BasicBlock> *NodeB = PDT.getNode(B);

		// Collect NodeA dominators set.
		SmallPtrSet<DomTreeNodeBase<BasicBlock>*, 16> NodeADoms;
		NodeADoms.insert(NodeA);
		DomTreeNodeBase<BasicBlock> *IDomA = NodeA->getIDom();
		while (IDomA) {
			NodeADoms.insert(IDomA);
			IDomA = IDomA->getIDom();
		}

		// Walk NodeB immediate dominators chain and find common dominator node.
		DomTreeNodeBase<BasicBlock> *IDomB = NodeB->getIDom();
		while (IDomB) {
			if (NodeADoms.count(IDomB) != 0)
				return IDomB->getBlock();

			IDomB = IDomB->getIDom();
		}

		return NULL;
	}

	bool OmitBranch::omit(BranchInst *branch) {
		BasicBlock *bb = branch->getParent();
		assert(branch->getNumSuccessors() > 0 && "The branch has no successor.");
		Function *func = bb->getParent();
		/* Calculate the nearest post dominator of <bb> */
		PostDominatorTree &PDT = getAnalysis<PostDominatorTree>(*func);
		BasicBlock *post_dominator_bb = branch->getSuccessor(0);
		for (unsigned i = 1; i < branch->getNumSuccessors(); i++) {
			/* findNearestCommonDominator does not work with a post dominator tree in LLVM 2.7 release */
			//post_dominator_bb = PDT.findNearestCommonDominator(post_dominator_bb, branch->getSuccessor(i));
			post_dominator_bb = find_nearest_common_post_dominator(PDT, post_dominator_bb, branch->getSuccessor(i));
			if (!post_dominator_bb)
				break;
		}
#if 0
		cerr << "bb = " << bb->getName().str() << endl;
		if (post_dominator_bb)
			cerr << "post dominator = " << post_dominator_bb->getName().str() << endl;
		else
			cerr << "no post dominator" << endl;
#endif
		/* Flood fill from <bb> until reaching <post_dominator_bb> */
		visited.clear();
		for (Function::iterator bi = func->begin(); bi != func->end(); ++bi)
			visited[bi] = false;
		for (unsigned i = 0; i < branch->getNumSuccessors(); i++)
			dfs(branch->getSuccessor(i), post_dominator_bb);
		/* If any visited BB has a sync operation, the branch cannot be omitted. */
		for (Function::iterator bi = func->begin(); bi != func->end(); ++bi) {
			if (!visited[bi])
				continue;
			// cerr << "Visited BB: " << bi->getNameStr() << endl;
			for (BasicBlock::iterator ii = bi->begin(); ii != bi->end(); ++ii) {
				if (is_call(ii) && !is_intrinsic_call(ii)) {
					vector<Function *> called_funcs = SF->get_called_functions(ii);
					for (size_t i = 0; i < called_funcs.size(); ++i) {
						// cerr << "Called function: " << called_funcs[i]->getName().str() << endl;
						if (SF->contains_sync_operation(called_funcs[i]))
							return false;
					}
				}
			}
		}
		return true;
	}

	char OmitBranch::ID = 0;
}

namespace {
	static RegisterPass<defens::OmitBranch> X(
			"omit-branch",
			"Determine whether a branch affects the trace of sync event",
			false,
			true); // is analysis
}

